package ResK.algorithms

import ResK.calculi.resolution._
import scala.collection._

object CottonSplit {
  private val rand = new scala.util.Random()

  def choosePivot (proof : Proof) : Atom = {
    val visitedProofs = new mutable.HashSet[Resolvent]
    val map = new mutable.HashMap[Atom,Int]
    var sum = 0
    def rec(p: Proof): Unit = p match {
      case _:Input => ()
      case resolvent:Resolvent =>
        if (!visitedProofs.contains(resolvent)) {
          rec(resolvent.left)
          rec(resolvent.right)
          val resolvent_additivity =
            (resolvent.clause.size -
             resolvent.left.clause.size max resolvent.right.clause.size) max 0
          sum += resolvent_additivity + 1
          val atom = resolvent.pivot._1.atom
          val atom_additivity =
            (if (map contains atom) map(atom) else 0) + resolvent_additivity + 1
          map += ((atom, atom_additivity))
          visitedProofs += resolvent
        }
    }
    rec(proof)

    // TODO : check map is not empty (raise an exception)

    // randomly pick an atom
    val it = map.keysIterator
    var ret = it.next()
    var r = rand.nextInt(sum) - map(ret)
    while (r > 0 && it.hasNext) {
      ret = it.next()
      r -= map(ret)
    }
    ret
  }

  def split (proof : Proof, pivot : Literal) : Proof = {
    // proof is supposed to be acyclic
    val computedProofs = new mutable.HashMap[Resolvent,Proof]
    def rec(p : Proof): Proof = p match {
      case i:Input => i
      case r:Resolvent => {
        if (!(computedProofs contains r)) {
          val pp = rec(r.left)
          val pn = rec(r.right)
          computedProofs += ((r,
            if (r.pivot._1 == pivot || !(pp.clause contains r.pivot._1)) pp else
            if (r.pivot._2 == pivot || !(pn.clause contains r.pivot._2)) pn else
            new Resolvent (pp, pn) ))
        }
        computedProofs(r)
      }
    }
    rec(proof)
  }

  def cottonSplit (proof : Proof, pivot : Atom) : Proof = {
    new Resolvent (split(proof, new Literal(pivot, true)),
                   split(proof, new Literal(pivot, false)))
  }

  def testCompression (filename : String, maxTimeSeconds : Long) : Float = {
    val begin_t = java.lang.System.currentTimeMillis
    var best_proof = ResK.calculi.ProofParser.getProofFromFile(filename)
    val initial_length = measures.length(best_proof)
    val reading_t = java.lang.System.currentTimeMillis
    println("Initial proof from " + filename + ": " + initial_length + " in " + (reading_t - begin_t) + "ms")

    while ((java.lang.System.currentTimeMillis - reading_t) < (maxTimeSeconds * 1000)) {
      val nproof = cottonSplit (best_proof, choosePivot(best_proof))
      val nlength = measures.length(nproof)
      if (nlength < measures.length(best_proof)) {
        println("Found proof: " + nlength + " in " + (java.lang.System.currentTimeMillis - reading_t) + "ms")
        best_proof = nproof
      }
    }

    println("Total time: " + (java.lang.System.currentTimeMillis - begin_t) + "ms")
    val compressionRatio = 1.f - (measures.length(best_proof).toFloat / initial_length.toFloat)
    println("Compression ratio for " + filename + ": " + compressionRatio)
    compressionRatio
  }

}
